from langchain.tools import BaseTool
from pydantic import BaseModel, Field, ConfigDict
from typing import Any, Type, Dict, List
from enum import Enum
import json
import random


# 1. 定义图表类型枚举（限制AI可选类型，减少错误）
class ChartType(str, Enum):
    LINE = "line"  # 折线图（适合趋势数据，如营收变化）
    BAR = "bar"    # 柱状图（适合对比数据，如各季度财报）
    PIE = "pie"    # 饼图（适合占比数据，如利润构成）


# 2. 定义工具输入参数模型（Pydantic 校验，确保必填参数存在）
class ShowInput(BaseModel):
    chart_type: ChartType = Field(...,
                                  description="图表类型，可选：line（折线图）、bar（柱状图）、pie（饼图），需匹配数据类型")
    data: Dict[str, Any] = Field(...,
                                 description="可视化数据，格式需匹配图表类型：\n"
                                             "- 折线/柱状图：{'x轴数据': ['2021','2022'], '系列1': [100,200], '系列2': [80,150]}\n"
                                             "- 饼图：{'数据': [{'name': '营收', 'value': 500}, {'name': '成本', 'value': 300}]}\n"
                                             "示例（柱状图）：{'x轴数据': ['2025Q1'], '营收(亿元)': [514.43], '净利润(亿元)': [268.47]}")
    title: str = Field(..., description="图表标题，需明确数据含义，如'贵州茅台2025年第一季度财报'")
    x_label: str = Field(None, description="x轴标签（仅折线/柱状图需要），如'季度'")
    y_label: str = Field(None, description="y轴标签（仅折线/柱状图需要），如'金额(亿元)'")


# 3. 可视化工具核心类（生成 ECharts 配置）
class ShowTool(BaseTool):
    model_config = ConfigDict(arbitrary_types_allowed=True)

    def __init__(self, **kwargs: Any):
        super().__init__(
            name="get_show_tool",  # AI 调用时的工具名
            description="用于财务数据可视化（如财报、营收），支持折线图/柱状图/饼图。当用户提'图表'/'可视化'/'画图'且有数据时，必须调用此工具。",
            **kwargs
        )

    args_schema: Type[BaseModel] = ShowInput  # 绑定输入参数模型

    def _run(
        self,
        chart_type: ChartType,
        data: Dict[str, Any],
        title: str,
        x_label: str = None,
        y_label: str = None
    ) -> str:
        """工具核心执行逻辑：根据参数生成 ECharts 配置 JSON"""
        try:
            # ========== 第一步：校验数据格式（按图表类型） ==========
            if chart_type in [ChartType.LINE, ChartType.BAR]:
                # 折线/柱状图：必须含“x轴数据”，且系列数据长度与 x 轴一致
                if "x轴数据" not in data:
                    raise ValueError("折线图/柱状图的 data 必须包含 'x轴数据' 字段（示例：{'x轴数据': ['2025Q1'], ...}）")
                x_len = len(data["x轴数据"])
                # 校验核心系列（营收、净利润、成本）长度
                for field in ["营收(亿元)", "净利润(亿元)", "成本(亿元)"]:
                    if field in data and len(data[field]) != x_len:
                        data[field] = data[field][:x_len]  # 长度不一致时截断
                        print(f"⚠️ 警告：{field} 数据长度与 x 轴不一致，已自动修正为前 {x_len} 个元素")

            elif chart_type == ChartType.PIE:
                # 饼图：必须含“数据”字段，且为 name-value 格式列表
                if "数据" not in data or not isinstance(data["数据"], list):
                    raise ValueError("饼图的 data 必须包含 '数据' 字段，且为列表（示例：{'数据': [{'name': '营收', 'value': 500}, ...]}）")


            # ========== 第二步：生成 ECharts 配置 ==========
            echarts_option = {
                "title": {
                    "text": title,           # 使用传入的标题
                    "left": "center",
                    "textStyle": {"fontSize": 16}
                },
                "tooltip": {
                    # 折线/柱状图用 axis 触发，饼图用 item 触发
                    "trigger": "axis" if chart_type in [ChartType.LINE, ChartType.BAR] else "item"
                },
                "legend": {
                    "data": [],  # 后续根据系列名填充
                    "bottom": 10
                },
                "toolbox": {  # 支持保存图片、重置图表
                    "show": True,
                    "feature": {"saveAsImage": {"show": True}, "restore": {"show": True}}
                },
                "series": []  # 后续根据数据填充系列
            }

            # 折线图/柱状图 特有配置
            if chart_type in [ChartType.LINE, ChartType.BAR]:
                echarts_option["xAxis"] = {
                    "type": "category",
                    "data": data["x轴数据"],
                    "name": x_label or "季度/年份",  # 无传入则用默认
                    "nameTextStyle": {"fontSize": 12}
                }
                echarts_option["yAxis"] = {
                    "type": "value",
                    "name": y_label or "金额（亿元）",  # 无传入则用默认
                    "nameTextStyle": {"fontSize": 12},
                    "min": 0  # y 轴从 0 开始，保证数据对比真实
                }
                # 提取所有系列名（排除 x 轴数据）
                series_names = [k for k in data.keys() if k != "x轴数据"]
                echarts_option["legend"]["data"] = series_names
                # 生成每个系列的配置
                for name in series_names:
                    echarts_option["series"].append({
                        "name": name,
                        "type": chart_type.value,
                        "data": data[name],
                        "itemStyle": {"color": f"#{random.randint(0, 0xFFFFFF):06x}"}  # 随机颜色区分系列
                    })

            # 饼图 特有配置
            elif chart_type == ChartType.PIE:
                echarts_option["series"].append({
                    "name": title,
                    "type": "pie",
                    "radius": ["40%", "70%"],  # 环形饼图（更美观）
                    "avoidLabelOverlap": False,
                    "itemStyle": {"borderColor": "#fff", "borderWidth": 2},  # 白色边框区分区块
                    "label": {"show": False, "position": "center"},  # 默认不显示标签（ hover 时强调显示）
                    "emphasis": {
                        "label": {"show": True, "fontSize": 16, "fontWeight": "bold"}
                    },
                    "labelLine": {"show": False},  # 隐藏标签连接线（更简洁）
                    "data": data["数据"]
                })
                # 图例显示所有数据项的 name
                echarts_option["legend"]["data"] = [item["name"] for item in data["数据"]]


            # ========== 第三步：转换为格式化 JSON 并返回 ==========
            return json.dumps(echarts_option, ensure_ascii=False, indent=2)

        except Exception as e:
            # 明确错误原因，便于调试
            return f"ERROR: ShowTool 生成失败 → {str(e)}（请检查参数格式是否符合要求）"

    def _arun(self, *args: Any, **kwargs: Any) -> None:
        """异步方法（若需异步调用，需在此实现，当前暂用同步逻辑）"""
        raise NotImplementedError("ShowTool 暂不支持异步调用")